#! /usr/bin/env python

#   SABBot: A Jabber Bot to control SABnzbd
#   Copyright (c) 2009 Syed Amer Gilani <syed.amer at gilani.eu>
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.

import warnings
warnings.simplefilter("ignore",DeprecationWarning)
import jabberbot
import urllib2
import json
import ConfigParser
import xmppony
import pycurl
import time
from xmppony.filetransfer import *
from urllib import urlencode
from os import remove
from os.path import exists, getsize
from threading import Thread

botcmd = jabberbot.botcmd

__author__ = 'Syed Amer Gilani <syed.amer at gilani.eu>'
__version__ = '0.1'

cfg_filename = 'SABBot.cfg'
jabber_user, jabber_pass, master, url, apikey = "", "", "", "", ""
 
class SABBot(jabberbot.JabberBot):
    def __init__(self, username, password):
        jabberbot.JabberBot.__init__(self, username, password)

    def connect_callback(self):
        self.conn.RegisterEventHandler(self.eventManager)
        SIFileTransfer().PlugIn(self.conn)
        self.conn.seekProxy()

    def callback_message(self, conn, mess):
        # only listen to my master
        if mess.getFrom().getStripped() != master:
            self.send(mess.getFrom(), "You are not my Master!", mess)
        else:
            return jabberbot.JabberBot.callback_message(self, conn, mess)

    def sab_query(self, query):
        query["apikey"] = apikey 
        response = urllib2.urlopen(urllib2.Request(url + "api?", urlencode(query)))
        response_data = response.read()
        return response_data
    
    @botcmd
    def status(self, mess, args):
        """Display Status Informations"""
        return_str = "\n"
        status_json = json.loads(self.sab_query({"mode":"qstatus", "output":"json"}))
        
        if status_json['paused']:
            return_str += "PAUSED\n"

        return_str += "Download speed: " + str(round(status_json['kbpersec'],2)) + " KB/s \n"
        return_str += "Queued: " + str(round(status_json['mbleft'],2)) + "/" + str(round(status_json['mb'],2)) + " MB" + "\n"
        return_str += "Download Dir: " + str(round(status_json['diskspace1'],2)) + " GB - Complete Dir: " + str(round(status_json['diskspace2'],2))  + " GB\n\n"

        jobs = status_json["jobs"]

        for job in jobs:
            return_str += job["filename"] + "\nRemain/Total " + str(round(job["mbleft"],2)) + "/" + str(round(job["mb"],2)) + " GB\n\n"
            
        self.send(mess.getFrom(), return_str, mess)

    @botcmd
    def pause(self, mess, args):
        """Pause the Queue"""
        return_str = self.sab_query({"mode":"pause"})
        self.send(mess.getFrom(), return_str, mess)

    @botcmd
    def resume(self, mess, args):
        """Resume the Queue"""
        return_str = self.sab_query({"mode":"resume"})
        self.send(mess.getFrom(), return_str, mess)

    @botcmd
    def addurl(self, mess, args):
        """Add nzb from URL"""
        query = {"mode":"addurl"}
        argl = args.split()
        if len(argl) == 0:
            self.send(mess.getFrom(), "No URL given", mess)    
            return
        else:
            query["name"] = argl[0]
            if len(argl) > 1:
                query["cat"] = argl[1]
                if len(argl) > 2:
                    query["pp"] = argl[2]
                    if len(argl) > 3:
                        query["script"] = argl[3]

        return_str = self.sab_query(query)
        self.send(mess.getFrom(), return_str, mess)

    @botcmd
    def speedlimit(self, mess, args):
        """Set speed limit in KB/s (0 to disable)"""
        query = {"mode":"config","name":"speedlimit"}
        argl = args.split()
        if len(argl) == 0:
            self.send(mess.getFrom(), "No Limit given", mess)    
            return
        else:
            query["value"] = argl[0]

        return_str = self.sab_query(query)
        self.send(mess.getFrom(), return_str, mess)

    @botcmd
    def exit( self, mess, args):
        """quit SABBot"""
        self.quit()

    def eventManager(self, realm, event, data):
        if realm == xmppony.filetransfer.NS_PROFILE_FT:
            if event == xmppony.si.SI_OFFERED:
                self.handleFileRequest(data)

    def fileupload(self, filename, size):
        timeout = time.time() + 60.0
        while True:
            if exists(filename):
                if getsize(filename) == int(size):
                    break
            if time.time() > timeout:
                return
            time.sleep(1)
          
        # since urlopen2 does not support file uploads i have to use curl here
        c = pycurl.Curl()
        c.setopt(c.POST, 1)
        c.setopt(c.URL, url + "api?apikey=" + apikey + "&mode=addfile")
        c.setopt(c.HTTPPOST, [("nzbfile", (c.FORM_FILE, str(filename)))])
        #c.setopt(c.VERBOSE, 1)
        c.perform()
        c.close()

        remove(filename)


    def handleFileRequest(self, info):
        self.conn.SIFileTransfer.AcceptStream(info['key'], info['filename'], cont=True)
        t = Thread(target=self.fileupload, args=(info['filename'],info['filesize'],))
        t.start()


def config():
    global jabber_user, jabber_pass, master, url, apikey
    if exists(cfg_filename):
        config = ConfigParser.ConfigParser()
        config.read(cfg_filename)

        if config.has_section('Jabber'):
            if config.has_option('Jabber','user') and config.has_option('Jabber','pass') and config.has_option('Jabber','master'):
                jabber_user = config.get('Jabber', 'user')
                jabber_pass = config.get('Jabber', 'pass')
                master = config.get('Jabber', 'master')
            else:
                print "No Jabber user/pass/master defined in Config"
                quit()            
        else:
            print "No [Jabber] Section in Config"
            quit()

        if config.has_section('SABnzbd'):
            if config.has_option('SABnzbd','url') and config.has_option('SABnzbd','apikey'):
                url = config.get('SABnzbd', 'url')
                apikey = config.get('SABnzbd', 'apikey')
            else:
                print "No Jabber url/apikey defined in Config"
                quit()            
        else:
            print "No [SABnzbd] Section in Config"
            quit()
        
    else:
        print "No Config File found"
        quit()

if __name__ == '__main__':
    config()
    bot = SABBot(jabber_user, jabber_pass)
    bot.serve_forever(bot.connect_callback)

